'''
v8 det 训练
使用方法
    1 准备数据

    2 训练
        cd yolov8
        python user_det_auto_train.py --root_path D:\data\231207huoni\trainV8Det_cable

        数据格式: D:\data\231207huoni\trainV8Det_cable（支持v8 det， v5det）
            _add_imgs/ 原始图片+芯歌格式标注
            format_data/ yolov8检测数据集（自动生成）（数据增强：随机裁剪；训练集为所有图片，验证集20%）
            models/ 存放模型（自动生成）
            predict/ 存放预测结果（自动生成）

'''

import argparse
import os.path

import yaml
from ultralytics import YOLO
from user_det_sg2yolo import sengo2yolo


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # parser.add_argument('--root_path', type=str, default=r'D:\data\240618anquandai_val\trainV8Det_v0624', help='root_path')

    parser.add_argument('--root_path', type=str, default=r'D:\data\240507fubatianxian\trainV8Det_all', help='root_path')
    # parser.add_argument('--root_path', type=str, default=r'D:\data\240507fubatianxian\trainV8Det_st5kakou')
    par = parser.parse_args()
    root_path = par.root_path

    # root_path = r'D:\data\231215安全带\trainV8Det_flball_blue'
    # root_path = r'/home/ps/zhangxiancai/data/231215anquandai/trainV8Det_flball_blue'

    # data
    if not os.path.exists(rf'{root_path}/format_data'):
        os.mkdir(rf'{root_path}/format_data')

    # s2y = sengo2yolo(glob_str=rf'{root_path}/_add_imgs',
    #                  yolov5_dir=rf'{root_path}/format_data')# sengo随机裁剪，转为v5格式数据集，生成data.yaml
    # s2y.run()
    data_yaml_path = f'{root_path}/format_data/data.yaml'

    # # train

    save_path = f'{root_path}/v5_v6.0_models'
    weights = f'yolov5/yolov5s_v6.0.pt'
    imgsz = 416
    device = '0'
    if not os.path.exists(save_path):
        os.mkdir(save_path)

    # weights = f'{root_path}/models/train/weights/last.pt'

    train_cmd_str = (f'python yolov5/train.py --epoch {100} --weights {weights} --data {data_yaml_path} --imgsz {imgsz} '
                     f'--batch-size {2} --project {save_path} '
                     f' --device {0}')
    print(train_cmd_str)
    os.system(train_cmd_str)